#include "Vm.h"
#include "Const.h"
#include "utility/Memory.h"
#include "Utils.h"
#include "Debug.h"
#include "Task.h"

// Virual Memory Management

uint8_t* bits;
MutexType mutex;
BitMapType bitmap;

// 这表示 textStart 是一个地址符号，它不是变量，也不是指针，只是一个“标签”，代表某个内存地址。
// 这种声明方式告诉编译器：“textStart 是一个地址，类型为 uint8_t 的数组的起始地址。”
// 实际上，不会分配任何数组空间，只是让你可以用 textStart 这个名字来引用链接脚本里的那个地址。
extern uint8_t textStart[], textEnd[], dataStart[], dataEnd[], kernelStart[];

static PDEType kernelPDT[PDE_COUNT] __attribute__((aligned(PAGE_SIZE)));

static PDEType* GetCurrentPDT(void)
{
    return (PDEType*)GetCurrentTask()->tss.cr3;
}

// Set bit map and return physical address
static uint32_t AllocatePage(BitMapType* bitmap, uint32_t count)
{
    LockMutex(&mutex);

    int addr = 0;
    int pageIndex = AllocateBitMap(bitmap, false, count);
    if (pageIndex >= 0)
    {
        addr = bitmap->start + pageIndex * bitmap->pageSize;
    }

    UnlockMutex(&mutex);

    return addr;
}

uint32_t AllocateKernelPage(void)
{
    return AllocatePage(&bitmap, 1);
}

/// @brief The FreeBitMap function releases a range of bits in a bitmap, 
/// starting at the specified address (addr) and spanning the given count (count). 
/// It ensures thread safety by locking a mutex during the operation, 
/// calculates the starting page index, and marks the specified range as free using SetBitMapBit.
/// @param bitmap 
/// @param addr 
/// @param count 
static void FreeBitMap(BitMapType* bitmap, uint32_t addr, uint32_t count)
{
    LockMutex(&mutex);

    int pageIndex = (addr - bitmap->start) / bitmap->pageSize;
    SetBitMapBit(bitmap, pageIndex, count, false);

    UnlockMutex(&mutex);
}

/// @brief 释放了Bitmap的位，找到了对应的PDE中的PTE，并释放了该PTE所对应的物理页
/// @param vaddr 
void FreeMemoryPage(uint32_t vaddr)
{
    // Kernel
    if (vaddr < TASK_ADDRESS_BASE)
    {
        FreeBitMap(&bitmap, vaddr, 1);
    }
    else
    {
        PTEType* pte = FindPTE(GetCurrentPDT(), vaddr, false);
        ASSERT(pte==(PTEType*)0 && pte->present);

        FreeBitMap(&bitmap, GetPhysicalPageAddress(pte), 1);
        pte->v= 0;
    }
}

void InitVM(MemoryInfoType* memoryInfo)
{
    InitMutex(&mutex);
    extern uint8_t memStart[];

    uint8_t* bitMapData = (uint8_t*)memStart;

    uint32_t allBytes;
    for (int i = 0; i < memoryInfo->memoryRegionCount; i++)
    {
        allBytes += memoryInfo->memoryRegionArray[i].size;
    }
    uint32_t freeSize = allBytes - KERNEL_END_ADDRESS;
    freeSize = DownAlign(freeSize, PAGE_SIZE);
    InitBitMap(&bitmap, bitMapData, KERNEL_END_ADDRESS, freeSize, PAGE_SIZE);
    bitMapData += GetBytesCount(freeSize / PAGE_SIZE);
    CreateKernelTable();
    SetPageDirectoryTable((uint32_t)kernelPDT);
}

void InitBitMap(BitMapType* bitmap, uint8_t* bits, uint32_t start, uint32_t size, uint32_t pageSize)
{
    InitMutex(bitmap->mutex);
    bitmap->count = size/ pageSize;
    bitmap->bitmap = bits;
    bitmap->pageSize = pageSize;
    bitmap->start = start;
    bitmap->size = size;
    Memset(bitmap->bitmap, 0, GetBytesCount(bitmap->count));
}

uint32_t GetBytesCount(uint32_t size)
{
    return (size+7)/8;
}

uint8_t GetBitMapBit(BitMapType* bitmap, uint32_t index)
{
    uint32_t byteIndex = index/8;
    uint32_t bitIndex = index%8;
    uint8_t bit = (bitmap->bitmap[byteIndex] >> bitIndex) & 1;
    return bit;
}

void SetBitMapBit(BitMapType* bitmap, uint32_t index, uint32_t count, bool flag)
{
    for (int i = 0; (i < count) && (index < bitmap->count); i++, index++)
    {
        bitmap->bitmap[index/8] |= (flag << (index%8));
    }
}

bool IsBitSet(BitMapType* bitmap, uint32_t index)
{
    return GetBitMapBit(bitmap, index) == 1;
}

/// @brief Allocate continous memory
/// @param bitmap
/// @param value
/// @param count
/// @return
int AllocateBitMap(BitMapType* bitmap, bool value, uint32_t count)
{
     int currentIndex = 0;
     int startIndex = 0;
     while(currentIndex < bitmap->count)
     {
         if (GetBitMapBit(bitmap, currentIndex) != value)
         {
            ++currentIndex;
            continue;
         }

         startIndex = currentIndex;

         int i;
         for (i = 0; (i < count)&&(currentIndex<bitmap->count); i++)
         {
            if (GetBitMapBit(bitmap, currentIndex++)!= value)
            {
                startIndex = -1;
                break;
            }
         }

         if (i>= count)
         {
            SetBitMapBit(bitmap, startIndex, count, !value);
            return startIndex;
         }
     }

     return -1;
}

/// @brief Find PTE
PTEType* FindPTE(PDEType* pdeTable, uint32_t vaddr, bool create)
{
    PTEType* pteTable;

    // Get pde in the pde table
    // Get pte table from pde
    PDEType* pde = pdeTable + GetPDEIndex(vaddr);
    if (pde->present)
    {
        pteTable = (PTEType*)GetPhysicalPageTableAddress(pde);
    }
    else
    {
        if (create == false)
        {
            return (PTEType*)0;
        }

        uint32_t pteAddress = AllocatePage(&bitmap, 1);
        if (pteAddress == 0)
        {
            return (PTEType*)0;
        }

        // pageTable address is already 4096 aligned
        if (vaddr >= TASK_ADDRESS_BASE)
        {
            pde->v = pteAddress | PDE_P | PTE_W | PDE_U;
        }
        else
        {
            pde->v = pteAddress | PDE_P | PTE_W;
        }
        
        pteTable = (PTEType*)pteAddress;

        Memset(pteTable, 0, PAGE_SIZE);
    }

    return pteTable + GetPTEIndex(vaddr);
}

/// @brief Create memory map and set pde and pte
/// @param pde
/// @param vaddr
/// @param paddr
/// @param count
/// @param perm
/// @return
int CreateMemoryMap(PDEType* pde, uint32_t vaddr, uint32_t paddr, uint32_t count, uint32_t perm)
{
    // How to get pte from virtual address?
    // 1. Get pde table which contains 1024 pde
    // 2. Get target pde from pde table by vaddr the high 10 bits
    // 3. Get pte table which contains 1024 pte
    // 4. Get target pte from pte table by vaddr the middle 10 bits
    // 5. Set pte permission and physical address

    for (int i = 0; i < count; i++)
    {
        PTEType* pte = FindPTE(pde, vaddr, true);
        if (pte == (PTEType*)0)
        {
            return -1;
        }

        ASSERT(pte->present == 0);

        pte->v = paddr | perm | PTE_P;

        vaddr += PAGE_SIZE;
        paddr += PAGE_SIZE;
    }
    return 0;
}

// MARK: BUG
/// @brief The AllocateMemory function allocates physical memory pages for a given virtual address range, sets bitmap
/// flags, and creates page directory and page table entries (PDE/PTE) to establish the memory mapping.
/// @param pdeAddress
/// @param vaddr
/// @param size
/// @param perm
/// @return
uint32_t AllocateMemory(uint32_t pdeAddress, uint32_t vaddr, uint32_t size, uint32_t perm)
{
    // Get virtual memory for building map
    uint32_t realVirtualAddress = DownAlign(vaddr, PAGE_SIZE);
    uint32_t realSize = UpAlign(size, PAGE_SIZE);
    // Get page count
    int pageCount = realSize / PAGE_SIZE;

    for (int i = 0; i < pageCount; i++)
    {
        // Get physical address
        // Set bitmap flag
        // Eg:
        //     Loop 1: 0x40000
        //     Loop 2: 0x42000
        uint32_t physicalAddress = AllocatePage(&bitmap, 1);
        // Build pde and pte map
        // Eg:
        //     Loop 1: PDE index 1 | PTE index 20 -> 0x40000
        //     Loop 2: PDE index 1 | PTE index 21 -> 0x42000
        int error = CreateMemoryMap((PDEType*)pdeAddress, vaddr, physicalAddress, 1, perm);
        if (error<0)
        {
            // Error
            FreeBitMap(&bitmap, realVirtualAddress, i+i);
            return -1;
        }
        vaddr+=PAGE_SIZE;
    }
    return 0;
}

uint32_t AllocateMemoryForCurrentTask(uint32_t vaddr, uint32_t size, int perm)
{
    return AllocateMemory(GetCurrentTask()->tss.cr3, vaddr, size, perm);
}

void CreateKernelTable(void)
{
    static MemoryMapType mapTable[] =
    {
        {kernelStart, textStart, 0, PTE_W}, // Kernel stack
        {textStart, textEnd, textStart, 0}, // Kernel code
        {dataStart, (void*)(MEMORY_EBDA_START-1), dataStart, PTE_W}, // Kernel data
        {(void*)KERNEL_END_ADDRESS, (void*)KERNEL_TASK_END, (void*)KERNEL_END_ADDRESS, PTE_W} // Task area
    };

    Memset(kernelPDT, 0, sizeof(kernelPDT));

    for (int i = 0; i < sizeof(mapTable) / sizeof(MemoryMapType); i++)
    {
        MemoryMapType* map = mapTable + i;

        uint32_t virtualStart = DownAlign((uint32_t)map->virtualStart, PAGE_SIZE);
        uint32_t virtualEnd = UpAlign((uint32_t)map->virtualEnd, PAGE_SIZE);
        uint32_t pageCount = (virtualEnd - virtualStart) / PAGE_SIZE;

        // FindPTE->AllocatePage->AllocateBitMap
        CreateMemoryMap(kernelPDT, virtualStart, map->physicalStart, pageCount, map->permission);
    }
}

/// @brief Create a pde table for a task and set bit map flag, 
/// meanwhile new pde table will copy kernel pde table data
/// @param
/// @return
int CreateTaskPageDirectoryTable(void)
{
    // Create pde table and allocate memory for this one
    PDEType* taskPDT = (PDEType*)AllocatePage(&bitmap, 1);
    if(taskPDT == (PDEType*)0)
    {
        return 0;
    }

    Memset((void*)taskPDT, 0, PAGE_SIZE);
    // uint32_t pdeStart
    uint32_t pdeIndex = GetPDEIndex(TASK_ADDRESS_BASE);
    for (int i = 0; i < pdeIndex; i++)
    {
        taskPDT[i].v = kernelPDT[i].v;
    }

    return (uint32_t)taskPDT;
}

/// @brief CopyUserVirtualMemoryAndPageTable is a function that duplicates the user-space virtual memory and page table from a given page directory, 
/// allocating new pages and copying their contents. If the operation fails at any point, 
/// it cleans up allocated resources and returns -1.
/// @param pageDir 
/// @return 
uint32_t CopyUserVirtualMemoryAndPageTable(uint32_t pageDir)
{
    uint32_t pageDirAddr = CreateTaskPageDirectoryTable();
    if (pageDirAddr == 0)
    {
        goto CopyFailed;
    }

    uint32_t pdeStartUser = GetPDEIndex(TASK_ADDRESS_BASE);
    PDEType* pde = (PDEType*)pageDir + pdeStartUser;
    for (int i = pdeStartUser; i < PDE_COUNT; i++, pde++)
    {
        if (!pde->present)
        {
            continue;
        }
        
        // Get physical address of pte table
        PTEType* pte = (PTEType*)GetPhysicalPageTableAddress(pde);
        for (int j = 0; j < PTE_COUNT; j++, pte++)
        {
            if (pte->present)
            {
                // MARK: BUG，超过128MB的地址会出错
                // Allocate memory for pte table
                uint32_t pageAddr = AllocatePage(&bitmap, 1);
                if (pageAddr==0)
                {
                    goto CopyFailed;
                }

                uint32_t vaddr = (i<<22) | (j<<12);
                int err = CreateMemoryMap((PDEType*)pageDirAddr, vaddr, pageAddr, 1, GetPTEPerm(pte));
                if (err<0)
                {
                    goto CopyFailed;
                }
                
                Memcpy((void*)pageAddr, (void*)vaddr, PAGE_SIZE);
            }
            else
            {
                continue;
            }
        }
    }

    return pageDirAddr;
    
CopyFailed:
    FreeUserVirtualMemoryAndPageTable(pageDirAddr);
    return -1;
}

void FreeUserVirtualMemoryAndPageTable(uint32_t pageDir)
{
    uint32_t pdeStartUser = GetPDEIndex(TASK_ADDRESS_BASE);
    PDEType* pde = (PDEType*)pageDir + pdeStartUser;

    for (int i = pdeStartUser; i < PDE_COUNT; i++, pde++)
    {
        if (!pde->present)
        {
            continue;
        }

        PTEType* pte = (PTEType*)GetPhysicalPageTableAddress(pde);
        for (int j = 0; j < PTE_COUNT; j++, pte++)
        {
            if (!pte->present)
            {
                continue;
            }
            // Free pte physical memory
            FreeBitMap(&bitmap, GetPhysicalPageAddress(pte), 1);
        }
        // Free pte table
        FreeBitMap(&bitmap, GetPhysicalPageTableAddress(pde), 1);
    }
    // Free pde table
    FreeBitMap(&bitmap, pageDir, 1);
}

int CopyDataFromOtherPageDir(uint32_t to, uint32_t pageDir, uint32_t from, uint32_t size)
{
    while (size>0)
    {
        uint32_t toPhyAddr = GetPhysicalAddress(pageDir, to);
        if (toPhyAddr==0)
        {
            return -1;
        }

        uint32_t offsetInPage = toPhyAddr & (PAGE_SIZE - 1);
        uint32_t currentSize = PAGE_SIZE - offsetInPage;
        if (currentSize>size)
        {
            currentSize = size;
        }
        
        Memcpy((void*)toPhyAddr, (void*)from, currentSize);
        
        size -= currentSize;
        to += currentSize;
        from += currentSize;
    }
    return 0;
}


/// @brief Get physical address through virtual address
/// @param pageDir 
/// @param vaddr 
/// @return 
uint32_t GetPhysicalAddress(uint32_t pageDir, uint32_t vaddr)
{
    PTEType* pte = FindPTE((PDEType*)pageDir, vaddr, false);
    if (pte == (PTEType*)0)
    {
        return 0;
    }
    return GetPhysicalPageAddress(pte) + (vaddr & (PAGE_SIZE-1));    
}

